import mne
import os
import math
import random
import numpy as np
import pandas as pd
import seaborn as sn
import matplotlib.pyplot as plt
from scipy.stats import f_oneway

#%%
# channel names https://github.com/TNTLFreiburg/braindecode/issues/44
#     Fz,
# FC3,FC1,FCz,FC2,FC4
# C5,C3,C1,Cz,C2,C4,C6
# CP3,CP1,CPz,CP2,CP4
# P1,Pz,P2
# POz
# 

# 'EEG-Fz':Fz       0
# 'EEG-0':FC3       1
# 'EEG-1':FC1       2    del
# 'EEG-2':FCz       3    del
# 'EEG-3':FC2       4    del
# 'EEG-4':FC4       5
# 'EEG-5':C5        6    del
# 'EEG-C3':C3       7
# 'EEG-6':C1        8    del
# 'EEG-Cz':Cz       9
# 'EEG-7':C2        10    del
# 'EEG-C4':C4       11
# 'EEG-8':C6        12    del
# 'EEG-9':Cp3       13
# 'EEG-10':Cp1      14    del
# 'EEG-11':Cpz      15    del
# 'EEG-12':Cp2      16    del
# 'EEG-13':Cp4      17
# 'EEG-14':P1       18    del
# 'EEG-Pz':Pz       19    del
# 'EEG-15':P2       20    del
# 'EEG-16':POz      21    del

# the rest channels: Fz,FC1,FC4,C3,Cz,C4,Cp3,Pz,Cp4
# do this because I want to test the algrithm on the 8 channels that same with my mii system

# when you want to delete the channels you can do:
# use_8ch = 1
    # if use_8ch == 1:
        # raw.drop_channels(['EEG-1', 'EEG-2','EEG-3','EEG-5','EEG-6',
        #                'EEG-7','EEG-8','EEG-10','EEG-11','EEG-12',
        #                'EEG-14','EEG-Pz','EEG-15', 'EEG-16'])
# do not use drop_channels again, the better way is : use np.delete to delete other channels, for example:
# data_8ch = np.delete(data, [2,3,4,6,8,10,12,14,15,16,18,19,20,21], axis=0)


#%%
# raw_data_folder = '/content/raw_data/'
# cleaned_data_folder = '/content/cleaned_data/'

# raw_data_folder = 'D:/materials/MI/competition/BCICIV_2a_gdf/'
# cleaned_data_folder = 'D:/materials/MI/competition/BCICIV_2a_gdf_cleaned/'
raw_data_folder = 'D:/materials/dataset/bci/BCICIV_2a_gdf/'
cleaned_data_folder = 'D:/materials/dataset/bci/BCICIV_2a_gdf_cleaned/'



files = os.listdir(raw_data_folder)

# Filtering out files with suffix 'E.gdf'
# filtered_files = [file for file in files if file.endswith('T.gdf')]

filtered_files = [file for file in files if file.endswith('01T.gdf')]

save_fif = 0
# new_file_path = os.path.join(cleaned_data_folder, 'All_Subjects.fif')
new_file_path = os.path.join(cleaned_data_folder, 'A01.fif')

# filePathName_save = cleaned_data_folder + "bcic_iv_2a_data_all_sub.p"
filePathName_save = cleaned_data_folder + "A01_filter4_100_fs250_t_5_25.p"


save_pickle = 1

#%%
raw_list = []

# Iterating through filtered files
for file in filtered_files:
    file_path = os.path.join(raw_data_folder, file)

    # Reading raw data
    raw = mne.io.read_raw_gdf(file_path, eog=['EOG-left', 'EOG-central', 'EOG-right'], preload=True)
    # Droping EOG channels
    raw.drop_channels(['EOG-left', 'EOG-central', 'EOG-right'])

    print(raw.ch_names)

    # High Pass Filtering 4-40 Hz
    raw.filter(l_freq=4, h_freq=100, method='iir')

    # Notch filter for Removal of Line Voltage
    raw.notch_filter(freqs=50)

    # Resampling Data
    # raw.resample(128, npad='auto')

    # Saving the modified raw data to a file with .fif suffix
    new_file_path = os.path.join(cleaned_data_folder, file[:-4] + '.fif')
    raw.save(new_file_path, overwrite=True)
    # Appending data to the list
    raw_list.append(raw)

final_raw = mne.concatenate_raws(raw_list)
if save_fif == 1:
    final_raw.save(new_file_path, overwrite=True)

#%%
events = mne.events_from_annotations(final_raw)
# events[1]
# {'1023': 1,
#  '1072': 2,
#  '276': 3,
#  '277': 4,
#  '32766': 5,
#  '768': 6,
#  '769': 7,
#  '770': 8,
#  '771': 9,
#  '772': 10}

# List of the events
# '1023': 1 Rejected trial
# '1072': 2 Eye movements
# '276': 3 Idling EEG (eyes open)
# '277': 4 Idling EEG (eyes closed)
# '32766': 5 Start of a new run
# '768': 6 Start of a trial
# '769': 7 Cue onset Left (class 1) : 0
# '770': 8 Cue onset Right (class 2) : 1
# '771': 9 Cue onset Foot (class 3) : 2
# '772': 10 Cue onset Tongue (class 4): 3
#%%

# Time choice:
# [0.5s, 2,5s] Post Cue on set: [3.75s, 5.75s]

epochs = mne.Epochs(final_raw, events[0], event_id=[7, 8, 9, 10], tmin=0.5, tmax=2.5, reject=None, baseline=None, preload=True)
data = epochs.get_data()
labels = epochs.events[:,-1]

print("Dataset's shape:",data.shape)

#%%
import pickle

bcic_iv_2a_data_all_sub = {}
bcic_iv_2a_data_all_sub['datas']=data
bcic_iv_2a_data_all_sub['labels']=labels
bcic_iv_2a_data_all_sub['fs']=250

# filePathName_save = 'D:\materials\MI\competition\BCICIV_2a_pickle/bcic_iv_2a_data_eval.p'

if save_pickle==1:
    pickle.dump(bcic_iv_2a_data_all_sub,open(filePathName_save,'wb'))

#%%
#%%
#%%
#%%

